#include "BufferView.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"
#include "../RenderGraph.h"

BufferView::Builder::Builder() noexcept = default;
BufferView::Builder::~Builder() noexcept = default;
BufferView::Builder::Builder(BufferView::Builder const& rhs) noexcept = default;
BufferView::Builder::Builder(BufferView::Builder && rhs) noexcept = default;
BufferView::Builder& BufferView::Builder::operator=(BufferView::Builder const& rhs) noexcept = default;
BufferView::Builder& BufferView::Builder::operator=(BufferView::Builder && rhs) noexcept = default;

BufferView::Builder& BufferView::Builder::SetMultiple(bool isMultiple) noexcept
{
	mImpl->isMultiple = isMultiple;
	return *this;
}

BufferView::Builder& BufferView::Builder::SetSize(VkDeviceSize size) noexcept
{
	mImpl->size = size;
	return *this;
}

BufferView::Builder& BufferView::Builder::SetUsage(VkBufferUsageFlags usage) noexcept
{
	// Have to contain VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT or VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT
	// Only Contain one
	if (usage & VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT)
	{
		assert(usage & VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT == 0);
	}
	else if (usage & VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT)
	{
		assert(usage & VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT == 0);
	}
	else
	{
		// None of two TEXEL_BUFFER_BIT
		assert(0);
	}

	mImpl->usage = usage;
	return *this;
}

BufferView::Builder& BufferView::Builder::SetSharingMode(VkSharingMode sharingMode) noexcept
{
	mImpl->sharingMode = sharingMode;
	return *this;
}

BufferView::Builder& BufferView::Builder::AddQueueFamily(VkQueueFlagBits queueFlag) noexcept
{
	mImpl->queueFamilys.push_back(queueFlag);
	return *this;
}

BufferView::Builder& BufferView::Builder::SetMemoryProperty(VkMemoryPropertyFlags memoryProperty) noexcept
{
	mImpl->memoryProperty = memoryProperty;
	return *this;
}

BufferView::Builder& BufferView::Builder::SetFormat(VkFormat format) noexcept
{
	mImpl->format = format;
	return *this;
}

BufferView::Builder& BufferView::Builder::SetOffset(VkDeviceSize offset) noexcept
{
	mImpl->offset = offset;
	return *this;
}

BufferView::Builder& BufferView::Builder::SetRange(VkDeviceSize range) noexcept
{
	mImpl->range = range;
	return *this;
}

BufferView::Builder& BufferView::Builder::SetData(void* data) noexcept
{
	mImpl->data = data;
	return *this;
}

RefCountPtr<BufferView> BufferView::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<BufferView>(NewObject<BufferView>(*mImpl));
}

RefCountPtr<BufferView> BufferView::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	BufferView* bufferView = NewObject<BufferView>(*mImpl);
	process->AddBuffer(name, bufferView);
	return RefCountPtr<BufferView>(bufferView);
}

void BufferView::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveBuffer(mName, this);
		}
		DeleteObject<BufferView>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveBuffer(mName, this);
			}
			DeleteObject<BufferView>(this);
			});
	}
}

BufferView::~BufferView()
{
	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		// vkDestroyBufferView(GVulkanRHI->GetLogicalDevice(), mBufferViews[i], nullptr);
		GRenderState->DelayGCBufferView(mBufferViews[i]);
	}
}

VkBufferView BufferView::GetBufferView(uint32 index) const
{
	if (mIsMultiple)
		return mBufferViews[0];
	if (index == UINT32_MAX)
		return mBufferViews[GetFrameIndex()];
	assert(index <= GetFrameCount());
	return mBufferViews[index];
}

BufferView::BufferView(BuilderDetails & builderDetails)
{
	mIsMultiple = builderDetails.isMultiple;

	mFlags = 0;
	mSize = builderDetails.size;
	mUsage = builderDetails.usage;
	mSharingMode = builderDetails.sharingMode;
	mQueueFamilys.assign(builderDetails.queueFamilys.begin(), builderDetails.queueFamilys.end());
	if (mQueueFamilys.size() > 1)
	{
		// multi queue
		assert(mSharingMode == VK_SHARING_MODE_CONCURRENT);
	}
	else
	{
		assert(mSharingMode == VK_SHARING_MODE_EXCLUSIVE);
	}
	mMemoryProperty = builderDetails.memoryProperty;

	mFormat = builderDetails.format;
	mOffset = builderDetails.offset;
	mRange = builderDetails.range;


	assert(mOffset < mSize);
	assert(mOffset % GVulkanRHI->GetDeviceProperties().limits.minTexelBufferOffsetAlignment == 0);

	if (mRange != VK_WHOLE_SIZE)
	{
		assert(mRange > 0);
		assert(mOffset + mRange < mSize);
		assert(mRange % GVkFormatSize[mFormat] == 0);
		assert(mRange / GVkFormatSize[mFormat] <= GVulkanRHI->GetDeviceProperties().limits.maxTexelBufferElements);
	}


	mBuffers.resize(GetFrameCount(), VK_NULL_HANDLE);
	mBufferMemorys.resize(GetFrameCount(), VK_NULL_HANDLE);
	mBufferViews.resize(GetFrameCount(), VK_NULL_HANDLE);

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		GVulkanRHI->CreateBuffer(mBuffers[i], mBufferMemorys[i], mSize, mUsage, mMemoryProperty, mSharingMode, mQueueFamilys);
	}

	// Decide use which way to write data
	if (builderDetails.data != nullptr)
	{
		// Only propertyFlags is VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT we can vkMapMemory
		// If propertyFlags is VK_MEMORY_PROPERTY_HOST_COHERENT_BIT we need to use vkFlushMappedMemoryRanges and vkInvalidateMappedMemoryRanges to write data
		if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT)
		{
			if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
			{
				for (uint32 i = 0; i < GetFrameCount(); ++i)
				{
					void* data;
					VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i], 0, mSize, 0, &data));
					memcpy(data, builderDetails.data, mSize);
					vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i]);
				}
			}
			else
			{
				for (uint32 i = 0; i < GetFrameCount(); ++i)
				{
					void* data;
					VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i], 0, mSize, 0, &data));
					memcpy(data, builderDetails.data, mSize);

					// Flush to make changes visible to the host
					VkMappedMemoryRange mappedMemoryRange;
					mappedMemoryRange.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE;
					mappedMemoryRange.memory = mBufferMemorys[i];
					mappedMemoryRange.offset = 0;
					mappedMemoryRange.size = mSize;
					vkFlushMappedMemoryRanges(GVulkanRHI->GetLogicalDevice(), 1, &mappedMemoryRange);
					vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i]);
				}
			}
		}
		else if (mMemoryProperty & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
		{
			// We can only use stage buffer to copy data on device
			// Usage need to set VK_BUFFER_USAGE_TRANSFER_DST_BIT
			assert(mUsage & VK_BUFFER_USAGE_TRANSFER_DST_BIT);

			VkBuffer stagingBuffer;
			VkDeviceMemory stagingMemory;

			GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, mSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

			void* data;
			VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, mSize, 0, &data));
			memcpy(data, builderDetails.data, mSize);
			vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

			VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();


			for (uint32 i = 0; i < GetFrameCount(); ++i)
			{
				VkBufferCopy bufferCopy{};
				bufferCopy.srcOffset = 0;
				bufferCopy.dstOffset = 0;
				bufferCopy.size = mSize;
				vkCmdCopyBuffer(commandBuffer, stagingBuffer, mBuffers[i], 1, &bufferCopy);
			}


			GRenderState->FlushSingleCommand(commandBuffer);

			vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
			vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
		}
		else
		{
			//TODO: more ways to copy data
			assert(0);
		}
	}


	VkBufferViewCreateInfo bufferViewCreateInfo{};
	bufferViewCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_VIEW_CREATE_INFO;
	bufferViewCreateInfo.format = mFormat;
	bufferViewCreateInfo.offset = mOffset;
	bufferViewCreateInfo.range = mRange;

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		bufferViewCreateInfo.buffer = mBuffers[i];
		VK_CHECK_RESULT(vkCreateBufferView(GVulkanRHI->GetLogicalDevice(), &bufferViewCreateInfo, nullptr, &mBufferViews[i]));
	}
}

